import argparse
import os
import numpy as np
# 导入自定义工具
from utils.h5data import read_h5_data
from utils.phy import *
START_TIME = '2024-05-30T00:00:00Z'

def plot_dvv_curve(te_bjt, dvv, dvv_err=None, fit_quality=None, output_file=None,title=''):
    '''
    绘制dv/v随时间变化的曲线
    '''
    from pylab import figure, rcParams, plt
    
    rcParams['font.family'] = 'Arial'
    rcParams['font.size'] = '8'
    
    fig = figure(figsize=(8, 2.5), dpi=500)
    
    # 主图：dv/v vs 时间, 原始单位： ms/s，所以是‰
    ax1 = plt.subplot(1, 1, 1)

    print(te_bjt[0],te_bjt[-1])
    plt.errorbar(te_bjt, dvv, yerr=dvv_err*3, 
                 fmt='.-', linewidth=0.7,color='tab:red', markersize=3, markeredgecolor='none',
                ecolor='tab:blue',elinewidth=0.6, capsize=1, capthick=0.6, barsabove=True
                     )
    plt.fill_between(te_bjt, dvv-dvv_err*3, dvv+dvv_err*3, alpha=0.4, color='tab:red')

    
    plt.xlabel('day (BJT)')
    plt.ylabel('dv/v(%)',color='tab:red')
    plt.title(title)
    plt.grid(True, alpha=0.3)
    plt.ylim([-np.std(dvv)*4,np.std(dvv)*4])
    
    # 子图：拟合质量(mse) vs 时间， 单位：ms
    PLOT_RMS  = False
    PLOT_TEMP = False
    PLOT_ST   = True # sensor temp
    PLOT_SP   = False
    PLOT_TIDE = False
    if PLOT_RMS:
        ax2 = plt.twinx(ax1)
        # plt.plot(te_bjt, fit_quality*100, '-', color='k',linewidth=0.6)
        # plt.ylabel('rms(%)', color='k')
        # plt.ylim(0, 200)
        # plt.yticks([0,50,100])

        plt.plot(te_bjt, fit_quality, '-', color='k',linewidth=0.6)
        plt.ylabel('rms(ms)', color='k')
        plt.ylim(0, np.std(fit_quality)*20)
        output_file = output_file.replace('.png', '.rms.png')
        # plt.yticks([0,50,100])
    # 子图：温度
    elif PLOT_TEMP:
        ax2 = plt.twinx(ax1)
        te_phy, phy = get_era5_T(START_TIME,key='t2m')  
        t2m= merge_data(te_phy,phy,t_merged=te_bjt)
        plt.plot(te_bjt, t2m, '--', color='blue',linewidth=0.6)
        plt.ylabel('T2m(deg)', color='blue')
        plt.ylim(t2m.min()-np.std(t2m), t2m.max()+np.std(t2m))
        output_file = output_file.replace('.png', '.temp.png')
    elif PLOT_ST:
        ax2 = plt.twinx(ax1)
        te_phy, phy = get_sensor_T(START_TIME, file = 'loc/P320_temp/P264.csv')
        t2m= merge_data(te_phy,phy,t_merged=te_bjt)
        plt.plot(te_bjt, t2m, '--', color='blue',linewidth=0.6)
        plt.ylabel('T_P264(deg)', color='blue')
        plt.ylim(t2m.min()-np.std(t2m), t2m.max()+np.std(t2m))
        output_file = output_file.replace('.png', '.st.png')

    elif PLOT_SP:
        ax2 = plt.twinx(ax1)
        te_phy, phy = get_era5_T(START_TIME,key='sp')  
        sp= merge_data(te_phy,phy,t_merged=te_bjt)/1000
        plt.plot(te_bjt, sp, '--', color='blue',linewidth=0.6)
        plt.ylabel('sp(kPa)', color='blue')
        plt.ylim(sp.min()-np.std(sp), sp.max()+np.std(sp))
        output_file = output_file.replace('.png', '.sp.png')
    elif PLOT_TIDE:
        ax2 = plt.twinx(ax1)
        te_phy, phy = get_solid(START_TIME, 30)
        solid= merge_data(te_phy,phy,t_merged=te_bjt)
        plt.plot(te_bjt, solid, '--', color='green',linewidth=0.6)
        plt.ylabel('solid', color='green')
        plt.ylim(solid.min()-np.std(solid), solid.max()+np.std(solid))
        output_file = output_file.replace('.png', '.tide.png')

    plt.tight_layout()
    
    if output_file:
        plt.savefig(output_file, dpi=600, bbox_inches='tight')
        print(f"Figure saved to {output_file}")
    
    return fig

def plot_dvv_vs_common( te_bjt, dvv, te_phy,phy,output_file, phy_name,title):
    
    ne = len(te_bjt)
    # 拟合过程
    t_shift = np.arange(-0.5,0.5,0.01)
    R = np.zeros(len(t_shift))
    phy_shifted = []
    for i,shift_i in  enumerate(t_shift):
        phy_i = merge_data(te_phy,phy,t_merged=te_bjt+shift_i)
        R[i] = np.sum((phy_i-phy_i.mean())*dvv)/ne
        phy_shifted.append(phy_i)
    R = np.abs(R)
    shift_best = t_shift[np.argmax(R)]
    phy_best = phy_shifted[np.argmax(R)]
    coffs = np.polyfit(phy_best,dvv,1)
    dvv_fit = np.polyval(coffs,phy_best)
    dvv_rms = dvv-dvv_fit

    # 决定系数 R²
    ss_res = np.sum(dvv_rms ** 2)
    ss_tot = np.sum((dvv - np.mean(dvv)) ** 2)
    R2 = 1 - (ss_res / ss_tot) if ss_tot != 0 else 0

    from pylab import figure, rcParams, plt
    rcParams['font.family'] = 'Arial'
    rcParams['font.size'] = '8'
    
    shift_best = t_shift[np.argmax(R)]

    # 图1 绘制搜索过程
    fig = figure(figsize=(3,2.5), dpi=500)
    plt.plot(t_shift*24,R, '-', color='blue',linewidth=0.6)
    plt.xlabel('time shift(hours)', color='blue')
    plt.ylabel('R', color='blue')
    output_fileNew = output_file.replace('.png', f'.{phy_name}.search.png')
    plt.savefig(output_fileNew, dpi=600, bbox_inches='tight')
    plt.close('all')

    # 图2 绘制最优拟合曲线的线性关系
    fig = figure(figsize=(3,3), dpi=500)
    plt.scatter(phy_best,dvv, marker='^', color='gray',
                linewidth=0.6, s=3, edgecolors='none')
    plt.scatter(phy_best,dvv_fit, marker='o', color='red',
                linewidth=0.6, s=3, edgecolors='none')
    plt.xlabel(f't2m(dt={shift_best*24:.1f}h)', color='blue')
    plt.ylabel(f'dvv(%)', color='blue')
    plt.title(f'R²={R2:.2f}')
    output_fileNew = output_file.replace('.png', f'.{phy_name}.best.png')
    plt.savefig(output_fileNew, dpi=600, bbox_inches='tight')
    plt.close('all')

    # 图3 绘制拟合效果
    fig = figure(figsize=(8, 2.5), dpi=500)
    ax1 = plt.subplot(1, 1, 1)
    print(te_bjt[0],te_bjt[-1])
    plt.errorbar(te_bjt, dvv_fit, yerr=dvv_err*3, 
                 fmt='.-', linewidth=0.7,color='k', markersize=3, markeredgecolor='none',
                ecolor='none',elinewidth=0.6, capsize=1, capthick=0.6, barsabove=True
                     )
    plt.errorbar(te_bjt, dvv, yerr=dvv_err*3, 
                 fmt='.-', linewidth=0.7,color='tab:red', markersize=3, markeredgecolor='none',
                ecolor='none',elinewidth=0.6, capsize=1, capthick=0.6, barsabove=True
                     )
    plt.fill_between(te_bjt, dvv-dvv_err*3, dvv+dvv_err*3, alpha=0.4, color='tab:red')
    plt.xlabel('day (BJT)')
    plt.ylabel('dv/v(%)',color='tab:red')
    plt.ylim([dvv.min()-np.std(dvv)*1,dvv.max()+np.std(dvv)])

    ax2 = plt.twinx(ax1)
    plt.plot(te_bjt, phy_best, '--', color='green',linewidth=0.6)
    plt.ylabel(f't2m(dt={shift_best*24:.1f}h)', color='green')
    plt.ylim([phy_best.min()-np.std(phy_best), phy_best.max()+np.std(phy_best)])

    plt.title(title)
    plt.grid(True, alpha=0.3)
    
    output_fileNew = output_file.replace('.png', f'.{phy_name}.fit.png')
    plt.savefig(output_fileNew, dpi=600, bbox_inches='tight')
    plt.close('all')

    return dvv_fit, phy_best,shift_best

def plot_dvv_VS_t2m(te_bjt, dvv, dvv_err=None, fit_quality=None, output_file=None,title=''):
    '''
    绘制dv/v随时间变化的曲线
    '''
    # 温度拟合
    
    te_phy, phy = get_era5_T(START_TIME,key='t2m')  
    

    dvv_fit, phy_best,shift_best = plot_dvv_vs_common(te_bjt, dvv, 
                                   te_phy,phy,output_file, 't2m',title=title)

    return dvv-dvv_fit
    
def plot_dvv_VS_tide(te_bjt, dvv, dvv_err=None, fit_quality=None, output_file=None,title=''):
    '''
    绘制dv/v随时间变化的曲线
    '''

    # 温度拟合
    
    te_phy, phy = get_solid(START_TIME, 30)

    dvv_fit, phy_best,shift_best = plot_dvv_vs_common(te_bjt, dvv, 
                                   te_phy,phy,output_file, 't2m',title=title)

    return dvv-dvv_fit

def plot_dvv_VS_st(te_bjt, dvv, dvv_err=None, fit_quality=None, output_file=None,title=''):
    '''
    绘制dv/v随时间变化的曲线
    '''

    # 温度拟合
    name = 'P264'
    te_phy, phy = get_sensor_T(START_TIME, file = f'loc/P320_temp/{name}.csv')

    dvv_fit, phy_best,shift_best = plot_dvv_vs_common(te_bjt, dvv, 
                                   te_phy,phy,output_file, name,title=title)

    return dvv-dvv_fit 


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Plot traceSeq h5 data')
    parser.add_argument('-input', default='', help='input h5 file')
    parser.add_argument('-figroot', default='figures/7.dt.change.figures', help='root to save figs')
    parser.add_argument('-vref', type=float, default=3000, help='reference velocity (m/s)')
    # parser.add_argument('-figroot', default='figures/debug', help='root to save figs')
    args = parser.parse_args()
    INPUT_FILE = args.input
    FIG_ROOT = args.figroot
    V_REF = args.vref

    datasets, args_in_file = read_h5_data(INPUT_FILE, 
                                         ['all_groups','MARKER','EMARKER','DATE'],
                                         group_name='metadata',
                                         read_attrs=True)
    all_groups = [i.decode() for i in datasets[0]]
    MARKER = datasets[1].decode() 
    EMARKER = datasets[2].decode() 
    DATE = datasets[3].decode() 

    FIG_ROOT = f'{FIG_ROOT}/7B.{DATE}.{EMARKER}'
    if not os.path.exists(FIG_ROOT):
        os.makedirs(FIG_ROOT)

    te, dvv, dvv_err, dt_err, fit_quality = get_dvv_data_from_7file(INPUT_FILE, v_ref=V_REF)
    te_bjt = te/3600/24 + 8/24

    mask = np.where(np.all([te_bjt<=20, te_bjt>=10],axis=0))
    # mask = np.where(np.all([te_bjt<=100]))
    dvv     = -dvv[mask]/10 # 0.001转换为百分比,并校正X的符号
    dvv_err = dvv_err[mask]/10
    dt_err  = dt_err[mask]
    te_bjt  = te_bjt[mask]
    print(te_bjt)

    # 绘制dv/v曲线
    dvv_fig = plot_dvv_curve(te_bjt, dvv, dvv_err, fit_quality,
                             title=f'{MARKER} DV/V vs Time',
                            output_file=f'{FIG_ROOT}/dvv_vs_time.{MARKER}.V{V_REF:06.1f}.png')
    dvv_rms = dvv

    dvv_rms = plot_dvv_VS_st(te_bjt, dvv_rms, dvv_err, fit_quality,
                             title=f'{MARKER} DV/V vs Time',
                            output_file=f'{FIG_ROOT}/dvv_vs_phy.{MARKER}.V{V_REF:06.1f}.png')

    # dvv_rms = plot_dvv_VS_t2m(te_bjt, dvv_rms, dvv_err, fit_quality,
    #                          title=f'{MARKER} DV/V vs Time',
    #                         output_file=f'{FIG_ROOT}/dvv_vs_phy.{MARKER}.V{V_REF:06.1f}.png')
    # # print(dvv_rms.mean())
    # # assert 1==2
    # dvv_rms = plot_dvv_VS_tide(te_bjt, dvv_rms, dvv_err, fit_quality,
    #                          title=f'{MARKER} DV/V vs Time',
    #                         output_file=f'{FIG_ROOT}/dvv_vs_phy.{MARKER}.V{V_REF:06.1f}.png')
